import cv2
import numpy as np
from matplotlib import pyplot as plt

def convert(img):
    b, g, r = cv2.split(img)
    dst = cv2.merge([r, g, b])
    return dst

img = cv2.imread("../images/xhd.jpg")
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

# 绿灯
lower_green = np.array([35, 100, 100])
upper_green = np.array([88, 255, 255])

# 黄色
lower_yellow = np.array([17, 100, 100])
upper_yellow = np.array([33, 255, 255])

# 红色
lower_red = np.array([0, 110, 145])
upper_red = np.array([11, 255, 255])

# 掩码
mask_g = cv2.inRange(hsv, lower_green, upper_green)
mask_y = cv2.inRange(hsv, lower_yellow, upper_yellow)
mask_r = cv2.inRange(hsv, lower_red, upper_red)

# 将掩码的背景改为白色，如果不改是紫色
mask_g_white_bg = cv2.bitwise_not(mask_g)  # 反转掩码
mask_y_white_bg = cv2.bitwise_not(mask_y)  # 反转掩码
mask_r_white_bg = cv2.bitwise_not(mask_r)  # 反转掩码

img = convert(img)

plt.subplot(2, 2, 1), plt.imshow(img), plt.title("Origin")
plt.subplot(2, 2, 2), plt.imshow(mask_r_white_bg, cmap='gray'), plt.title('Red')
plt.subplot(2, 2, 3), plt.imshow(mask_y_white_bg, cmap='gray'), plt.title("Yellow")
plt.subplot(2, 2, 4), plt.imshow(mask_g_white_bg, cmap='gray'), plt.title("Green")
plt.show()